# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: leon

import abc
import datetime
from collections import namedtuple
import os
from typing import Union

import pytz

from coin.base.datetime_util import (to_datetime, to_timestamp_int)
from coin.base.hostname import get_hostname
from coin.base.timestamp import get_timestamp
from coin.flow.archive import ensure_dir_exists
from coin.flow.topic_record import TopicRecordReader, TopicRecordWriter, CrcLengthError
from coin.proto.coin_request_pb2 import (AccountRequestProto, StrategyRequestProto)

LogReaderReturn = namedtuple('LogReaderReturn', ['timestamp', 'log_content'])

LogReaderTailReturn = namedtuple('LogReaderTailReturn', ['file_name', 'timestamp', 'log_content'])


def to_date_str(trading_date):
  return trading_date.strftime("%Y%m%d")


def to_time_str(timestamp):
  return datetime.datetime.utcfromtimestamp(timestamp / 1e9).strftime("%Y%m%d-%H%M%S")


class ProtoLogProvider(object):
  def __init__(self):
    self._writer = None

    self._request_id = 0
    self._response_id = 0

    self._log_reset_timestamp = None

    self._write_callback = None

  def _log_reset_timestamp_updated(self, log_ts):
    reset = False
    log_time = to_datetime(log_ts)
    next_date = log_time.date() + datetime.timedelta(days=1)
    next_dt = datetime.datetime(
        next_date.year, next_date.month, next_date.day).replace(tzinfo=pytz.UTC)
    if self._log_reset_timestamp is None or log_ts >= self._log_reset_timestamp:
      reset = True
      self._log_reset_timestamp = to_timestamp_int(next_dt)
    return reset

  @abc.abstractmethod
  def gen_writer(self, timestamp):
    raise NotImplementedError()

  def update_writer(self, timestamp):
    if self._log_reset_timestamp_updated(timestamp):
      writer = self.gen_writer(timestamp)
      self.set_writer(writer)

  def set_writer(self, writer):
    self._writer = writer

  def set_write_callback(self, callback):
    self._write_callback = callback

  def get_next_request_id(self):
    self._request_id += 1
    return self._request_id

  def get_next_response_id(self):
    self._response_id += 1
    return self._response_id

  def write(self, protobuf, flush: bool = True, timestamp=None):
    timestamp = timestamp or get_timestamp()
    self.update_writer(timestamp)
    if self._write_callback is not None:
      self._write_callback(protobuf)
    self._writer.write(None, 0, None, timestamp, protobuf.SerializeToString())
    if flush:
      self.flush()

  def flush(self):
    if self._writer is not None:
      self._writer.flush()

  def close(self):
    if self._writer is not None:
      self._writer.close()
      self._writer = None


class ProtoLogReader(object):
  def __init__(self, paths):
    self._readers = []
    self._curr_reader_idx = -1
    self._prepare_readers(paths)

  def _prepare_readers(self, paths):
    for path in paths:
      self._readers.append(TopicRecordReader(path))
    if len(self._readers) > 0:
      self._curr_reader_idx = 0

  def read(self):
    if self._curr_reader_idx == -1:
      return None
    record = None
    try:
      record = self._readers[self._curr_reader_idx].read()
    except CrcLengthError:
      pass
    if record:
      return LogReaderReturn(record.timestamp, record.data)
    elif self._curr_reader_idx < len(self._readers) - 1:
      self._curr_reader_idx += 1
      return self.read()
    else:
      self._curr_reader_idx = -1
      return None


def gen_log_reader_iter(paths):
  log_reader = ProtoLogReader(paths)
  while True:
    ret = log_reader.read()
    if ret is not None:
      yield ret
    else:
      break


class ProtoLogReaderTail(object):
  def __init__(self, log_dir):
    self._log_dir = log_dir
    self._file_reader_map = {}
    self._add_new_logs(log_dir)

  def _add_new_logs(self, log_dir):
    if not os.path.isdir(log_dir):
      return
    for file_name in os.listdir(log_dir):
      if not file_name in self._file_reader_map:
        self._file_reader_map[file_name] = \
            TopicRecordReader(os.path.join(log_dir, file_name))

  def read(self):
    for file_name, reader in self._file_reader_map.items():
      record = reader.read()
      if record:
        return LogReaderTailReturn(file_name, record.timestamp, record.data)
    self._add_new_logs(self._log_dir)
    return LogReaderTailReturn(None, None, None)


def gen_log_reader_iter_tail(trading_date: datetime.datetime, log_root, machine=None):
  date_str = to_date_str(trading_date)
  machine = machine or get_hostname()
  log_dir = os.path.join(log_root, machine, date_str)
  log_reader = ProtoLogReaderTail(log_dir)
  while True:
    yield log_reader.read()
